package com.loyalove.demo.akka

import akka.actor.{Actor, ActorSystem, Props}
import com.typesafe.config.ConfigFactory

import scala.collection.mutable
import scala.concurrent.duration._

class SparkMaster extends Actor {

  // 保存workerInfo
  val workers = mutable.HashMap[String, WorkerInfo]()

  override def receive: Receive = {

    case RegisterWorkerInfo(workId, core, ram) =>
      if (!workers.contains(workId)) {
        val workerInfo = WorkerInfo(workId, core, ram, System.currentTimeMillis())

        workers += ((workId, workerInfo))

        println(s"------ $workId: $core|$ram registered ------")

        sender() ! RegisteredWorkerInfo(workId, core, ram)
      }

    case Heartbeat(workId) =>
      if(workers.contains(workId)){
        val workerInfo = workers(workId)
        workerInfo.time = System.currentTimeMillis()
      }

    case CheckTimeoutWorkerTask =>
      import context.dispatcher
      context.system.scheduler
        .schedule(0 millis, 6000 millis, self, CheckTimeoutWorker)

    case CheckTimeoutWorker =>
      val size = workers.size
      val currTime = System.currentTimeMillis()
      workers.values
        .filter(worker => currTime - worker.time > 3000)
        .foreach(worker => {
          workers.remove(worker.workId)
        })
      val finalSize = workers.size
      println(s"------ workers status: $finalSize alive, ${size - finalSize} removed ------")
  }
}


object SparkMaster {

  def main(args: Array[String]): Unit = {

    if (args.length != 3) {
      println(
        s"""
           |请输入参数： <host> <port> <masterName>
        """.stripMargin)

      sys.exit()
    }

    val host = args(0)
    val port = args(1)
    val masterName = args(2)

    val config = ConfigFactory.parseString(
      s"""
         |akka.actor.provider=akka.remote.RemoteActorRefProvider
         |akka.remote.netty.tcp.hostname=$host
         |akka.remote.netty.tcp.port=$port
         """.stripMargin)
    val actorSystem = ActorSystem("sparkMaster", config)

    // 创建自己的actorRef对象
    val masterActorRef = actorSystem.actorOf(Props[SparkMaster], masterName)

    // 启动检查超时worker任务
    masterActorRef ! CheckTimeoutWorkerTask
  }

}
